/*********************************************************************************
* Copyright 2019 Huawei Technologies Co.,Ltd.
* Licensed under the Apache License, Version 2.0 (the "License"); you may not use
* this file except in compliance with the License.  You may obtain a copy of the
* License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software distributed
* under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
* CONDITIONS OF ANY KIND, either express or implied.  See the License for the
* specific language governing permissions and limitations under the License.
**********************************************************************************
*/
#include <ctype.h>
#include <stdio.h>
#include <stdlib.h>
#include <sys/stat.h>
#include <sys/types.h>
#include <time.h>
#include <getopt.h>
#include <limits.h>

#if defined __GNUC__ || defined LINUX
#include <string.h>
#include <getopt.h>
#include <strings.h>
#include <unistd.h>
#include <pthread.h>
#else
#include "getopt.h"
#endif

#include "eSDKOBS.h"
#include "obs_demo_common.h"
#include "securec.h"
#include "enclave_proxy.h"

#include "kms.h"
#include "qtsm_lib.h"
#include "attestation.h"
#define  CIPHERTEXT_FILE "./text"
#define  PLAINTEXT_LENGTH 1000
#define  PARENT_CID 3
#define  ENCLAVE_DEMO_BIN "enclave_demo_bin"
#define  OBS_MAX_COUNT 10
#define  ARGS_NUM 11
#define  VSOCK_PORT_SHIFT 2
#define  PLAINTEXT_KEY_LENGTH 128

static int vsock_port;
static char plaintext_key[PLAINTEXT_KEY_LENGTH];

struct ObsObject {
    char *key;
    unsigned char *buffer;
    unsigned int buffer_size;
    int is_get;
    int vsock_cid;
    int vsock_port;
};

struct ProcessedData {
    char *data;
    int data_len;
};

struct KmsInfo {
    int vsock_cid;
    int vsock_port;
    char *endpoint;
    char *project_id;
    char *key_id;
};

static int test_get_object(char *key)
{
    char *file_name = "./test";
    obs_object_info object_info;
    obs_options option;
    option.bucket_options.host_name = HOST_NAME;
    option.bucket_options.bucket_name = BUCKET_NAME;
    option.bucket_options.access_key = ACCESS_KEY_ID;
    option.bucket_options.secret_access_key = SECRET_ACCESS_KEY;
    get_object_callback_data data;
    int rc;

    init_obs_options(&option);
    memset_s(&object_info, sizeof(object_info), 0, sizeof(obs_object_info));
    object_info.key =key;

    data.ret_status = OBS_STATUS_BUTT;
    data.outfile = write_to_file(file_name);

    obs_get_conditions getcondition;
    memset_s(&getcondition, sizeof(getcondition), 0, sizeof(obs_get_conditions));
    init_get_properties(&getcondition);
    // The starting position of the reading
    getcondition.start_byte = 0;
    // Read length, default is 0: read to the end of the object
    obs_get_object_handler get_object_handler = {
        {&get_properties_callback, &get_object_complete_callback},
        &get_object_data_callback
    };

    get_object(&option, &object_info, &getcondition, 0, &get_object_handler, &data);
    if (OBS_STATUS_OK == data.ret_status) {
        printf("get %s object successfully.\n", BUCKET_NAME);
        rc = 0;
    } else {
        printf("get %s object faied(%s).\n", BUCKET_NAME, obs_get_status_name(data.ret_status));
        rc = -1;
    }
    fclose(data.outfile);
    return rc;
}

static int test_put_object_from_buffer(struct ObsObject *obs_object)
{
    // Initialize option
    obs_options option;
    init_obs_options(&option);
    option.bucket_options.host_name = HOST_NAME;
    option.bucket_options.bucket_name = BUCKET_NAME;
    option.bucket_options.access_key = ACCESS_KEY_ID;
    option.bucket_options.secret_access_key = SECRET_ACCESS_KEY;
    int rc;

    // Initialize upload object properties
    obs_put_properties put_properties;
    init_put_properties(&put_properties);

    // Initialize the structure that stores the uploaded data
    put_buffer_object_callback_data data;
    memset_s(&data, sizeof(data), 0, sizeof(put_buffer_object_callback_data));
    // Assign buffer to the uploaded data structure
    data.put_buffer = obs_object->buffer;
    // set buffersize
    data.buffer_size = obs_object->buffer_size;

    // Set callback function
    obs_put_object_handler putobjectHandler = {
        { &response_properties_callback, &put_buffer_complete_callback},
          &put_buffer_data_callback
    };

    put_object(&option, obs_object->key, obs_object->buffer_size, &put_properties, 0, &putobjectHandler, &data);
    if (OBS_STATUS_OK == data.ret_status) {
        printf("put object %s from buffer successfully.\n", obs_object->key);
        rc = 0;
    } else {
        printf("put object %s from buffer failed(%s).\n", obs_object->key, obs_get_status_name(data.ret_status));
        rc = -1;
    }
    return rc;
}

static int test_interact_with_obs(struct ObsObject *obs_object)
{
    struct connect_info conn = {0};
    int rc;

    memset_s(&conn, sizeof(conn), 0, sizeof(struct connect_info));
    /*------ obs init------*/
    obs_initialize(OBS_INIT_ALL);
    set_online_request_max_count(OBS_MAX_COUNT);
    /* setup proxy */
    rc = setup_proxy(&conn, obs_object->vsock_cid, obs_object->vsock_port);
    if (rc != PX_NO_ERROR) {
        printf("setup_proxy obs error. rc is %d\n", rc);
        goto err_conn;
    }

    if (obs_object->is_get & 1) {
        rc = test_get_object(obs_object->key);
    } else {
        rc = test_put_object_from_buffer(obs_object);
    }

err_conn:
    close_proxy(&conn);
    obs_deinitialize();
    return rc;
}

static int read_ciphertext_from_file(struct ProcessedData *ciphertext_data)
{
    FILE *fp;
    int rc;
    int len;

    fp = fopen(CIPHERTEXT_FILE, "r");
    if (fp == NULL) {
        printf("reading test error\n");
        return -1;
    }

    fseek(fp, 0, SEEK_END);
    len = ftell(fp);
    if (len < 0 || len > INT_MAX) {
        printf("invalid ciphertext");
        rc = -1;
        goto fclose_fp;
    }
    ciphertext_data->data_len = len;
    rewind(fp);
    ciphertext_data->data = (char *)malloc((ciphertext_data->data_len + 1) * sizeof(char));
    if (!ciphertext_data->data) {
        printf("malloc ciphertext_data error\n");
        rc = -1;
        goto fclose_fp;
    }
    memset_s(ciphertext_data->data, (ciphertext_data->data_len + 1) * sizeof(char), 0,
            (ciphertext_data->data_len + 1) * sizeof(char));
    rc = fread(ciphertext_data->data, 1, ciphertext_data->data_len * sizeof(char), fp);
    if (rc != len) {
        printf("fread ciphertext from ciphertext error\n");
        goto free_ciphertext_data_data;
    }
    printf("%s\n", ciphertext_data->data);
    fclose(fp);
    return 0;
free_ciphertext_data_data:
    if (ciphertext_data->data) {
        free(ciphertext_data->data);
        ciphertext_data->data = NULL;
    }
fclose_fp:
    fclose(fp);
    return rc;
}

static unsigned long test_kms_decrypt_data(sig_params_t params, keyid_handle_t handle,
                                           const struct ProcessedData *ciphertext, struct ProcessedData *plaintext)
{
    plain_cipher_buff_t data_buffer = {0};
    unsigned long status;
    unsigned int i;
    data_buffer.data_in = ciphertext->data;
    data_buffer.data_in_len = ciphertext->data_len;
    data_buffer.data_out = plaintext->data;
    data_buffer.data_out_len = &(plaintext->data_len);

    printf("\n***********kms input **************\n");
    printf("data_in: %s \n data_in_len : %d\n data_out: %s\n data_out_len:%d\n",
           data_buffer.data_in, data_buffer.data_in_len, data_buffer.data_out, data_buffer.data_out_len);
    printf("************kms input end ************\n");
    /* call kms_decrypt_data api */
    status = kms_decrypt_data_blocking(&params, &handle, &data_buffer);
    if (status == KMS_SUCCESS) {
        printf("plain_data_len is %u\n", *(data_buffer.data_out_len));
        for (i = 0; i < *(data_buffer.data_out_len); i++)
            printf("%c", data_buffer.data_out[i]);
        printf("\n");
    }
    return status;
}

static int test_kms_decrypt(struct KmsInfo *kms_info, struct ProcessedData *plaintext_data)
{
    int rc;
    struct connect_info conn;
    sig_params_t params;
    keyid_handle_t handle;
    unsigned long status;
    struct ProcessedData *ciphertext_data;

    /* get ciphertext from text file */
    ciphertext_data = malloc(sizeof(struct ProcessedData));
    if (!ciphertext_data) {
        printf("alloc memory for ciphertext_data error. \n");
        return -1;
    }
    rc = read_ciphertext_from_file(ciphertext_data);
    if (rc < 0) {
        printf("read ciphertext from file error in test_kms_decrpyt. \n");
        goto free_ciphertext_data;
    }

    /* decrypt the ciphertext invoking kms decrypt api */
    rc = NO_ERROR;
    sig_params_init(&params);
    sig_str_t ak = sig_str(ACCESS_KEY_ID);
    sig_str_t sk = sig_str(SECRET_ACCESS_KEY);
    sig_str_t host = sig_str(kms_info->endpoint);
    sig_str_t uri_prefix = sig_str(kms_info->project_id);
    /* Set sign params */
    params.key = ak;
    params.secret = sk;
    params.host = host;
    params.uri_prefix = uri_prefix;
    /* Set key_id */
    handle.key_id = kms_info->key_id;
    handle.len = strlen(kms_info->key_id);

    /* setup proxy */
    memset_s(&conn, sizeof(conn), 0, sizeof(struct connect_info));
    rc = setup_proxy(&conn, kms_info->vsock_cid, kms_info->vsock_port);
    if (rc != PX_NO_ERROR) {
        printf("setup_proxy in kms decryption error.\n");
        rc = -1;
        goto free_ciphertext_data;
    }

    /* test decrypt data api */
    status = test_kms_decrypt_data(params, handle, ciphertext_data, plaintext_data);
    if (status != KMS_SUCCESS) {
        printf("invoking kms decrypt api err: %ld.\n", status);
        rc = -1;
    }
close_conn:
    close_proxy(&conn);
free_ciphertext_data:
    free(ciphertext_data);
    return rc;
}

static void help()
{
    printf("%s [-p vsock_port] [-a ak] [-s sk] \n " \
           "[-c ciphertext_key] [-l plaintext_key] [-b bucket_name] \n" \
           "[-k key_id] [-m kms_endpoint] [-o obs_endpoint] [-j project_id] \n", ENCLAVE_DEMO_BIN);
}

static int get_ciphertext(struct ObsObject *obs_object)
{
    int rc;

    obs_object->is_get = 1;
    obs_object->vsock_cid = PARENT_CID;
    obs_object->vsock_port = vsock_port;
    rc = test_interact_with_obs(obs_object);
    if (rc < 0) {
        printf("get ciphertext from obs error\n");
        return -1;
    }
    return 0;
}

static int invoke_kms_decryption(struct KmsInfo *kms_info, struct ProcessedData *plaintext_data)
{
    int rc;

    kms_info->vsock_cid = PARENT_CID;
    kms_info->vsock_port = vsock_port + 1;
    rc = test_kms_decrypt(kms_info, plaintext_data);
    if (rc < 0) {
        printf("kms_decrpt error\n");
    }
    return rc;
}

static int process_input(int argc, char **argv, struct KmsInfo *kms_info, struct ObsObject *obs_object)
{
    int rc;
    int c;
    const struct option options[] = {
        {"help", no_argument, 0, 'h'},
        {"vsock_port", required_argument, 0, 'p'},
        {"ak", required_argument, 0, 'a'},
        {"sk", required_argument, 0, 's'},
        {"ciphertext_key", required_argument, 0, 'c'},
        {"plaintext_key", required_argument, 0, 'l'},
        {"bucket_name", required_argument, 0, 'b'},
        {"key_id", required_argument, 0, 'k'},
        {"kms_endpoint", required_argument, 0, 'm'},
        {"obs_endpoint", required_argument, 0, 'o'},
        {"project_id", required_argument, 0, 'j'},
        {0, 0, 0, 0}
    };

    while ((c = getopt_long(argc, argv, "hc:p:a:s:c:l:b:k:m:o:j:", options, NULL)) != -1) {
        switch (c) {
            case 'h':
            case '?':
                help();
                return 0;
            case 'p':
                vsock_port = atoll(optarg);
                break;
            case 'a':
                rc = strcpy_s(ACCESS_KEY_ID, sizeof(ACCESS_KEY_ID), optarg);
                break;
            case 's':
                rc = strcpy_s(SECRET_ACCESS_KEY, sizeof(SECRET_ACCESS_KEY), optarg);
                break;
            case 'c':
                obs_object->key = optarg;
                break;
            case 'l':
                rc = strcpy_s(plaintext_key,  sizeof(plaintext_key), optarg);
                break;
            case 'b':
                rc = strcpy_s(BUCKET_NAME, sizeof(BUCKET_NAME), optarg);
                break;
            case 'k':
                kms_info->key_id = optarg;
                break;
            case 'm':
                kms_info->endpoint = optarg;
                break;
            case 'o':
                rc = strcpy_s(HOST_NAME, sizeof(HOST_NAME), optarg);
                break;
            case 'j':
                kms_info->project_id = optarg;
                break;
            default:
                printf("unsupport option %c\n", c);
                break;
        }
        if (rc < 0) {
            printf("get parameters fails\n");
            return -1;
        }
    }
    return 0;
}

static int put_data_to_obs(struct ObsObject *obs_object, struct ProcessedData *plaintext_data)
{
    int rc;

    obs_object->key = plaintext_key;
    obs_object->is_get = 0;
    obs_object->buffer_size = plaintext_data->data_len;
    obs_object->buffer = plaintext_data->data;
    obs_object->vsock_port = vsock_port + VSOCK_PORT_SHIFT;
    rc = test_interact_with_obs(obs_object);
    if (rc < 0)
        printf("put data to obs err\n");
    return rc;
}

int main(int argc, char **argv)
{
    struct ObsObject obs_object;
    int rc;
    struct KmsInfo kms_info;
    struct ProcessedData plaintext_data;

    /*---------process input-------*/
    rc = process_input(argc, argv, &kms_info, &obs_object);
    if (rc < 0)
        return -1;

    /*--------get ciphertext----*/
    rc = get_ciphertext(&obs_object);
    if (rc < 0)
        return -1;

    /*---- invoke kms decryption api -----*/
    plaintext_data.data_len = PLAINTEXT_LENGTH;
    plaintext_data.data = (char *)malloc(plaintext_data.data_len * sizeof(char));
    memset_s(plaintext_data.data, plaintext_data.data_len * sizeof(char),
             0, plaintext_data.data_len * sizeof(char));
    rc = invoke_kms_decryption(&kms_info, &plaintext_data);
    if (rc < 0)
        goto free_plaintext_data;

    /*------ put plaintext to obs------*/
    rc = put_data_to_obs(&obs_object, &plaintext_data);
    if (rc < 0)
        printf("put plaintext to obs error\n");

free_plaintext_data:
    if (plaintext_data.data) {
        free(plaintext_data.data);
        plaintext_data.data = NULL;
    }
    return rc;
}

